import torch.nn as nn
import torch.nn.functional as F

class DPCNN(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_classes, num_filters, kernel_size=3):
        super(DPCNN, self).__init__()
        
        # 词嵌入层
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        
        # 区域嵌入层
        self.region_embedding = nn.Conv1d(embed_dim, num_filters, kernel_size, padding=1)
        
        # 残差块
        self.conv_block = nn.Sequential(
            nn.Conv1d(num_filters, num_filters, kernel_size, padding=1),
            nn.ReLU(),
            nn.Dropout(0.5),  # 添加 Dropout
            nn.Conv1d(num_filters, num_filters, kernel_size, padding=1),
            nn.ReLU()
        )
        
        # 全连接层
        self.fc = nn.Linear(num_filters, num_classes)

    def forward(self, x):
        # 词嵌入并调整维度
        x = self.embedding(x).permute(0, 2, 1)
    
        # 区域嵌入
        x = F.relu(self.region_embedding(x))
    
        # 残差块
        residual = x
        x = self.conv_block(x)
        x = x + residual  # 残差连接
    
        # 全局池化
        x = F.adaptive_avg_pool1d(x, 1).squeeze(2)
        x = F.relu(x)
        out = self.fc(x)
        return out

